import warnings

import pandas as pd
from reco.datasets import loadMovieLens100k
from reco.metrics import rmse
from reco.recommender import FunkSVD


class FunkSVDPredict:
    warnings.filterwarnings('ignore')

    def __init__(self, data: pd.DataFrame):
        self.model = None
        self.data = data

    def fitModel(self, k=64, learning_rate=0.006,
                 regularizer=0.2, iterations=5, method='stochastic',
                 bias=True) -> None:
        model = FunkSVD(k=k,
                        learning_rate=learning_rate,
                        regularizer=regularizer,
                        iterations=iterations,
                        method=method,
                        bias=bias)
        model.fit(X=self.data,
                  formatizer={'user': 'userId', 'item': 'movieId', 'value': 'rating'},
                  verbose=True)
        self.model = model

    def getID(self, userID, listLength=20, maxLength=1000) -> list:
        l1, l2 = [], []
        for i in range(maxLength):
            l1.append(userID)
            l2.append(i + 1)
        test = pd.DataFrame({
            "userId": l1,
            "itemId": l2
        })
        l2 = self.model.predict(test, formatizer={'user': 'userId', 'item': 'itemId'})
        _ids = []
        for i in range(listLength):
            _ids.append(l2.index(max(l2)) + 1)
            l2.remove(max(l2))
        return _ids


# 读取数据 创建模型
model = FunkSVDPredict(pd.read_csv("./dataset/movielens100k.csv"))
# 训练模型
model.fitModel()
# 返回推荐列表
print(model.getID(1))
# [5, 9, 13, 23, 31]




